import os
import numpy as np
import  yaml

LANDMARK_FILE_NOT_EXIST = 901
PREDICTION_FOLDER_NOT_EXIST = 902
GT_FOLDER_NOT_EXIST = 903
LANDMARK_POINTS_NUMBER_ERROR = 904
CONFIGURE_FILE_NOT_EXIST = 905
CONFIGURE_FILE_WRONG_FORMAT = 906

def get_landmark_from_landmark_file(landmark_file,  landmark_number=68):
    '''

    :param landmark_file:
    :param file_exts:
    :param landmark_number:
    :return:
    '''
    landmark_pts = np.zeros([landmark_number, 2])
    if not os.path.exists(landmark_file):
        print('LANDMARK_FILE_NOT_EXIST')
        return LANDMARK_FILE_NOT_EXIST
    with open(landmark_file,'r') as infile:
        lines = infile.readlines()
        for num in range(len(lines-1)):
            index = num - 3
            if index < 0:
                continue
            line = lines[index]
            line = line.strip()

            x = line.split(' ')[0]
            y = line.split(' ')[1]
            landmark_pts[index,0] = x
            landmark_pts[index,1] = y
    return landmark_pts


def compute_error(pre_folder, gt_folder, file_exts='pts', landmark_number=68):
    if not os.path.exists(pre_folder):
        print('the landmark prediction folder does not exist')
        return PREDICTION_FOLDER_NOT_EXIST
    if not os.path.exists(gt_folder):
        print('the landmark gt folder does not exist')
        return GT_FOLDER_NOT_EXIST
    landmark_files = os.listdir(pre_folder)

    landmark_error_list = []
    for landmark_file in landmark_files:
        file_ext = os.path.basename(landmark_file).split('.')[1]
        if file_ext != file_exts:
            print('landmark file extension must be pts')
            continue
        pre_file = os.path.join(pre_folder, landmark_file)
        gt_file = os.path.join(gt_folder, landmark_file)
        landmark_pre = get_landmark_from_landmark_file(pre_file, landmark_number)
        landmark_gt = get_landmark_from_landmark_file(gt_file, landmark_number)

        if landmark_number == 68:
            interocular_distance = np.linalg.norm(landmark_gt[37, :]-landmark_gt[46, :])
        elif landmark_number == 51:
            interocular_distance = np.linalg.norm(landmark_gt[20, :]-landmark_gt[29, :])
        else:
            print(' support landmark points number 68 or 51')
            return LANDMARK_POINTS_NUMBER_ERROR
        sum = 0
        for point_index in range(len(landmark_number)):

            sum = sum + np.linalg.norm(landmark_pre[point_index, :]-landmark_gt[point_index, :])

        error_per_image = sum / (landmark_number * interocular_distance)
        landmark_error_list.append(error_per_image)
    assert len(landmark_error_list) != 0
    return landmark_error_list


def caculate_landmark_error_weight_sum(pre_folder, gt_folder, file_exts='pts', landmark_number=68):
    landmark_error_list = compute_error(pre_folder, gt_folder, file_exts, landmark_number)
    landmark_error_set = set(landmark_error_list)
    error_sum = 0.0
    images_num = len(landmark_error_list)
    for error_item in landmark_error_set:
        error_weight = landmark_error_list.count(error_item)
        error_sum = error_sum + error_weight*error_item
    error_sum = error_sum/images_num
    return error_sum

def main( default_configure_path='./configure.yaml' ):
    if not os.path.exists(default_configure_path):
        print('./configure.yaml does not exist')
        return CONFIGURE_FILE_NOT_EXIST
    config_file = open(default_configure_path)
    config = yaml.load(config_file)
    if not isinstance(config, dict):
        print('wrong configure file format of ./configure.yaml')
        return CONFIGURE_FILE_WRONG_FORMAT
    if 'pre_folder' not in config \
            or 'gt_folder' not in config\
            or 'landmark_number' not in config\
            or 'landmark_file_extension' not in config:
        print("./configure.yaml format is like:"
              "pre_folder : 'path' "
              "gt_folder : 'other path' "
              'landmark_number : 68 '
              "landmark_file_extension : 'pts' " )
        return CONFIGURE_FILE_WRONG_FORMAT
    pre_folder = config['pre_folder']
    gt_folder = config['gt_folder']
    landmark_number = config['landmark_number']
    file_exts = config['landmark_file_extension']
    error_value = caculate_landmark_error_weight_sum(pre_folder, gt_folder, file_exts, landmark_number)
    print('---------------- landmark error value is---------------')
    print(error_value)

# if __name__=='__main__':
#     main()
